"""
通过pytorch完成手写体识别的训练和测试
"""
# ①　导入torch相关的包
import torch as pt
import torchvision as ptv
from torch.utils.data import TensorDataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
import os
import sys

# ②　通过torchvision导入mnist数据集，并分为训练集和验证集
from torchvision.datasets.mnist import MNIST

# arguments
BATCH_SIZE = 4
N_EPOCHS = 2

# ③　将数据集传入DataLoader
data_dir = '../../../../../large_data/DL2/pt/mnist/'
if not os.path.exists(data_dir):
    print(f'Dir {data_dir} is not existed!')
    sys.exit(1)
train = MNIST(data_dir, True, ptv.transforms.ToTensor, download=False)
test = MNIST(data_dir, False, ptv.transforms.ToTensor, download=False)
x_train = pt.unsqueeze(train.data, dim=1) / 255.
x_test = pt.unsqueeze(test.data, dim=1) / 255.
y_train = train.targets
y_test = test.targets
print('x_train', x_train.size())
print('y_train', y_train.size())
print('x_test', x_test.size())
print('y_test', y_test.size())
ds_train = TensorDataset(x_train, y_train)
ds_test = TensorDataset(x_test, y_test)
dl_train = DataLoader(ds_train, BATCH_SIZE, shuffle=True)
dl_test = DataLoader(ds_test, BATCH_SIZE, shuffle=True)


# ④　创建一个类，构造函数中定义一个ModuleList
class ModuleList(pt.nn.Module):

    # ⑤　包含卷积层，池化层，卷积层，池化层，droupout层，全连接层
    def __init__(self):
        super().__init__()
        self.conv1 = pt.nn.Conv2d(1, 32, (3, 3), (1, 1), (1, 1))
        self.bn1 = pt.nn.BatchNorm2d(32)
        self.relu1 = pt.nn.ReLU()
        self.pool1 = pt.nn.AvgPool2d((2, 2), (2, 2))
        self.conv2 = pt.nn.Conv2d(32, 64, (3, 3), (1, 1), (1, 1))
        self.bn2 = pt.nn.BatchNorm2d(64)
        self.relu2 = pt.nn.ReLU()
        self.pool2 = pt.nn.AvgPool2d((2, 2), (2, 2))
        self.dropout = pt.nn.Dropout(0.3)
        self.fc = pt.nn.Linear(7*7*64, 10)
        self.softmax = pt.nn.LogSoftmax(dim=1)  # ATTENTION dim=1

    # ⑥　定义前向传播
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.reshape(-1, 7*7*64)
        x = self.dropout(x)
        x = self.fc(x)
        x = self.softmax(x)
        return x


# ⑦　创建模型对象
model = ModuleList()

# ⑧　将数据放在cuda上
device = 'cuda:0' if pt.cuda.is_available() else 'cpu'
print('device:', device)
device = pt.device(device)

# ⑨　将模型放在cuda上
model = model.to(device)

# ⑩　定义损失函数
loss = pt.nn.NLLLoss()


def accuracy(y_true, y_pred):
    y_true = y_true.long()
    y_pred = y_pred.argmax(dim=1)
    acc = pt.eq(y_true, y_pred).float().mean()
    return acc


# 11　定义优化器
optim = pt.optim.Adam(model.parameters(), lr=0.001)


# 12　开始训练
def process_data(dl, label, is_train, epoch):
    avg_loss, avg_acc = 0., 0.
    dl_len = len(dl)
    group = int(np.ceil(dl_len / 10))
    for i, (bx, by) in enumerate(dl):
        bx = bx.float().to(device)
        by = by.long().to(device)
        if is_train:
            model.train(True)
            optim.zero_grad()
            pred = model(bx)
            lossv = loss(pred, by)
            lossv.backward()
            optim.step()
            model.train(False)
        else:
            model.train(False)
            pred = model(bx)
            lossv = loss(pred, by)
        accv = accuracy(by, pred)
        lossv = lossv.detach().cpu().numpy()
        accv = accv.detach().cpu().numpy()
        avg_loss += lossv
        avg_acc += accv
        if i % group == 0:
            print(f'{label}: epoch{epoch}: #{i + 1}: loss = {lossv}, acc = {accv}')
    if i % group != 0:
        print(f'{label}: epoch{epoch}: #{i + 1}: loss = {lossv}, acc = {accv}')
    avg_loss /= i + 1
    avg_acc /= i + 1
    return avg_loss, avg_acc


loss_his, acc_his, loss_his_val, acc_his_val = [], [], [], []
for epoch in range(N_EPOCHS):
    lossv, accv = process_data(dl_train, 'train', True, epoch + 1)
    print(f'epoch#{epoch + 1}: avg_loss = {lossv}, avg_acc = {accv}')
    loss_his.append(lossv)
    acc_his.append(accv)
    lossv_val, accv_val = process_data(dl_test, 'val', False, epoch + 1)
    print(f'epoch#{epoch + 1}: avg_loss = {lossv}, avg_acc = {accv}, avg_loss_val = {lossv_val}, avg_acc_val = {accv_val}')
    loss_his_val.append(lossv_val)
    acc_his_val.append(accv_val)

# 13　进行测试
print('Testing ...')
lossv_test, accv_test = process_data(dl_test, 'test', False, 1)
print(f'Test loss = {lossv_test}, acc = {accv_test}')
print('Tested')

# plotting
spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])

spn += 1
plt.subplot(spr, spc, spn)
plt.title('loss')
plt.plot(loss_his, label='train')
plt.plot(loss_his_val, label='val')
plt.legend()

spn += 1
plt.subplot(spr, spc, spn)
plt.title('acc')
plt.plot(acc_his, label='train')
plt.plot(acc_his_val, label='val')
plt.legend()

plt.show()
